'''
Created on 2009-8-28

@author: lz
'''

import pefile
from lib_family import SampleFamily, get_pickle_familys
import lib_decrypt

def judge_family(pe, family):
    bHitMajor = False
    bHitImports = False
    bKillImports = False
    bHitDetails = False
    bKillDetails = False

    # match Major first
    for section in pe.sections:
        if section.data.find(family.Major) != -1:
            bHitMajor = True
            break
    
    if not bHitMajor:
        return False
    
    # match import
    imports = list(family.Imports)
    if imports == ['', '', '']:
        bKillImports = True
    else:
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            for imp in entry.imports:
                if imp.name in imports:
                    imports.remove(imp.name)
        if imports == []:
            bHitImports = True
    
    # match detail
    if family.Signs == ('', '', ''):
        bKillDetails = True
    else:
        for section in pe.sections:
            for item in family.Signs:
                if item!= '' and section.data.find(item) != -1:
                    bHitDetails = True
                    break
    
    if bKillImports and bKillDetails:
        return True
    elif bKillImports:
        return bHitDetails
    elif bKillDetails:
        return bHitImports
    else:
        if family.Method == 'AND':
            return bHitDetails and bHitImports
        else:
            return bHitImports or bHitDetails

def judgesample(filename):

    pe = pefile.PE(filename)
    
    family_list = get_pickle_familys()
    for family in family_list:
        if judge_family(pe, family):
            print family.name
            if family.Fix == 'URL':
                fun = getattr(lib_decrypt, 'decrypt_%s' % (family.name,))
                url = fun(filename)
                print url
            break

if __name__ == '__main__':
    filename = r'e:\xp_share\2exJW3dsaTgWrf5uAPadmHN_unpacked.dll'
    judgesample(filename)
    print 'done'